#include <atomic>
#include <chrono>
#include <iostream>
#include <thread>
#include <vector>

#include "data_dumper.h"

constexpr int kNumThreads = 10;
constexpr int kNumPushesPerThread = 10000;
constexpr int kSleepTimeUs = 10;

std::atomic<int> total_pushed{0};

void PushData(DataDumper& dumper) {
  for (int i = 0; i < kNumPushesPerThread; ++i) {
    std::string op_name =
        "Default/"
        "TestXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
        "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX/"
        "MatMul_Op_" +
        std::to_string(i);
    auto report_data =
        std::make_unique<OpRangeData>(0, 0, 0, 0, 0, 0, 0, false, op_name);

    // report_data->input_dtypes.push_back("test_dtype1");
    // report_data->input_dtypes.push_back("test_dtype2");
    // report_data->input_dtypes.push_back("test_dtype3");

    // report_data->input_shapes.push_back({1, 2, 3});
    // report_data->input_shapes.push_back({4, 5, 6});
    // report_data->input_shapes.push_back({7, 8, 9});

    // report_data->stack.push_back("test_stack1.py/test_func1/test_line1");
    // report_data->stack.push_back("test_stack2.py/test_func2/test_line2");
    // report_data->stack.push_back("test_stack3.py/test_func3/test_line3");

    // report_data->module_hierarchy.push_back("test_module1");
    // report_data->module_hierarchy.push_back("test_module2");
    // report_data->module_hierarchy.push_back("test_module3");
    dumper.Report(std::move(report_data));
    total_pushed.fetch_add(1, std::memory_order_relaxed);
    std::this_thread::sleep_for(std::chrono::microseconds(kSleepTimeUs));
  }
}

int main() {
  auto start_time = std::chrono::high_resolution_clock::now();
  DataDumper& dumper = DataDumper::GetInstance();
  dumper.Init("./tmp", 10 * 1024);
  dumper.Start();

  std::vector<std::thread> threads;
  for (int i = 0; i < kNumThreads; ++i) {
    threads.emplace_back(PushData, std::ref(dumper));
  }

  for (auto& t : threads) {
    t.join();
  }

  dumper.Stop();

  // 检查是否有数据丢失
  int expected_total = kNumThreads * kNumPushesPerThread;
  int actual_total = total_pushed.load(std::memory_order_relaxed);
  std::cout << "thread num: " << kNumThreads
            << ", push num per thread: " << kNumPushesPerThread << std::endl;
  std::cout << "expected total: " << expected_total << std::endl;
  std::cout << "actual total: " << actual_total << std::endl;

  auto end_time = std::chrono::high_resolution_clock::now();
  std::cout << "cost time: "
            << std::chrono::duration_cast<std::chrono::milliseconds>(end_time -
                                                                     start_time)
                   .count()
            << "ms" << std::endl;

  return 0;
}